{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "# Copyright (c) Meta Platforms, Inc. and affiliates.\n",
    "#\n",
    "# This source code is licensed under the MIT license found in the\n",
    "# LICENSE file in the root directory of this source tree.\n",
    "import argparse\n",
    "import dataclasses\n",
    "import random\n",
    "import sys\n",
    "import timeit\n",
    "from typing import Tuple\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "\n",
    "# from home_robot.mapping.voxel import SparseVoxelMap\n",
    "from home_robot.utils.point_cloud_torch import unproject_masked_depth_to_xyz_coordinates\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from home_robot.datasets.scannet import ScanNetDataset, ScanNetModalities\n",
    "from home_robot.datasets.scannet.scannet_constants import SCANNET_COLOR_MAP_300\n",
    "colors = list(SCANNET_COLOR_MAP_300.values())\n",
    "random.shuffle(colors)\n",
    "\n",
    "\n",
    "data = ScanNetDataset(\n",
    "    root_dir = '/private/home/ssax/home-robot/src/home_robot/home_robot/datasets/scannet/data',\n",
    "    frame_skip = 180,\n",
    "    n_classes = 50,\n",
    "    # referit3d_config = ReferIt3dDataConfig(),\n",
    "    # scanrefer_config = ScanReferDataConfig(),\n",
    ")\n",
    "\n",
    "# Load specific scene\n",
    "# scene0192_00 -- small scene\n",
    "# 'scene0000_00' -- large scene\n",
    "idx = data.scene_list.index(\"scene0000_00\") #'scene0000_00'\n",
    "# idx = 0\n",
    "print(f\"Loaded images of (h: {data.height}, w: {data.width}) - resized from ({data.DEFAULT_HEIGHT},{data.DEFAULT_WIDTH})\")\n",
    "scene_obs = data.__getitem__(idx, show_progress=True)\n",
    "\n",
    "# # Load GT mesh\n",
    "# from pytorch3d.io import IO, load_obj, load_ply\n",
    "# from pytorch3d.structures import Pointclouds\n",
    "# def transform_points(points, transform_mat):\n",
    "#     return (torch.cat([points, torch.ones_like(points[:,:1])], dim=-1) @ transform_mat.T)[...,:3]\n",
    "\n",
    "# scene_id = scene_obs['scan_name']\n",
    "# print(\"Loading GT mesh for\", scene_id)\n",
    "# pointcloud = IO().load_pointcloud(data.root_dir / f'scans/{scene_id}/{scene_id}_vh_clean.ply')\n",
    "# verts = transform_points(pointcloud.points_packed(), scene_obs['axis_align_mats'][0])\n",
    "# pointcloud_aligned =  Pointclouds(points = [verts], features = [pointcloud.features_packed()])\n",
    "# print(f\"GT mesh: {len(verts)} verts\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(scene_obs['instance_2ds'][0].max())\n",
    "plt.imshow(scene_obs['instance_2ds'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "import cv2\n",
    "import numpy as np\n",
    "import torch\n",
    "from detectron2.structures import BoxMode, Instances, Boxes\n",
    "from detectron2.utils.visualizer import Visualizer, ColorMode, _create_text_labels, GenericMask\n",
    "from detectron2.data import MetadataCatalog\n",
    "\n",
    "def show_instance_image(rgb, instance_map, class_list, class_id_to_name, scale=1.0, custom_class_colors=None, custom_draw_fn = None):\n",
    "    # original_image: numpy array with shape (H, W, 3)\n",
    "    # instance_image: torch.long tensor with shape (H, W)\n",
    "    # class_list: list of ints representing classes\n",
    "    \n",
    "    assert rgb.max() <= 1.0\n",
    "    original_image = (rgb.cpu().numpy() * 255).astype(np.uint8)  # replace with your image path or method to load image\n",
    "    instance_image_np = instance_map.cpu().numpy()  # converting torch tensor to numpy array\n",
    "\n",
    "    # Extracting unique instances and their counts\n",
    "    instances_ids = np.unique(instance_image_np)\n",
    "\n",
    "    # Creating masks and bounding boxes for each instance\n",
    "    masks = [instance_image_np == i for i in instances_ids if i != -1]\n",
    "    boxes = [cv2.boundingRect(np.uint8(mask)) for mask in masks]  # returns (x,y,w,h)\n",
    "\n",
    "    # Convert (x,y,w,h) to (x1, y1, x2, y2) format\n",
    "    boxes = [(x, y, x+w, y+h) for x,y,w,h in boxes]\n",
    "\n",
    "    # Creating a dict for detected instances\n",
    "    height, width = original_image.shape[:2]\n",
    "    instances_obj = Instances((height, width))\n",
    "    instances_obj.pred_boxes = Boxes(boxes)\n",
    "    instances_obj.pred_masks = masks\n",
    "    if class_list is not None:\n",
    "        instances_obj.pred_classes = torch.as_tensor(class_list, dtype=torch.int64)\n",
    "        # Prepare metadata for visualization\n",
    "        # Assuming COCO dataset categories for example. Replace with your own if different.\n",
    "    else:\n",
    "        instances_obj.pred_classes = torch.as_tensor(instances_ids, dtype=torch.int64)\n",
    "    \n",
    "    if class_id_to_name is not None:\n",
    "        class_names = list(class_id_to_name.values())\n",
    "    else:\n",
    "        class_names = {i: str(i) for i in instances_ids}\n",
    "\n",
    "    # Clear the existing metadata if it's safe to do so\n",
    "    dataset_name = \"segmenter_data\"\n",
    "    if dataset_name in MetadataCatalog:\n",
    "        MetadataCatalog.remove(dataset_name)\n",
    "    # Create (or get) a MetadataCatalog for your dataset and set its thing_classes attribute        \n",
    "    metadata = MetadataCatalog.get(dataset_name)\n",
    "    metadata.set(thing_classes=class_names)\n",
    "    if custom_class_colors is not None:\n",
    "        metadata.thing_colors = custom_class_colors\n",
    "\n",
    "\n",
    "    # Visualize the results\n",
    "    v = Visualizer(original_image[:, :, ::-1], metadata=metadata, scale=scale)\n",
    "    if custom_draw_fn is None:\n",
    "        out = v.draw_instance_predictions(instances_obj)\n",
    "    else:\n",
    "        out = custom_draw_fn(v, instances_obj)\n",
    "    visualized_image = out.get_image()[:, :, ::-1]  # convert from RGB to BGR\n",
    "    return visualized_image\n",
    "\n",
    "\n",
    "def draw_instance_predictions(self, predictions):\n",
    "    \"\"\"\n",
    "    Draw instance-level prediction results on an image.\n",
    "\n",
    "    Args:\n",
    "        predictions (Instances): the output of an instance detection/segmentation\n",
    "            model. Following fields will be used to draw:\n",
    "            \"pred_boxes\", \"pred_classes\", \"scores\", \"pred_masks\" (or \"pred_masks_rle\").\n",
    "\n",
    "    Returns:\n",
    "        output (VisImage): image object with visualizations.\n",
    "    \"\"\"\n",
    "    self._instance_mode = ColorMode.SEGMENTATION\n",
    "    boxes = predictions.pred_boxes if predictions.has(\"pred_boxes\") else None\n",
    "    scores = predictions.scores if predictions.has(\"scores\") else None\n",
    "    classes = predictions.pred_classes.tolist() if predictions.has(\"pred_classes\") else None\n",
    "    labels = _create_text_labels(classes, scores, self.metadata.get(\"thing_classes\", None))\n",
    "    keypoints = predictions.pred_keypoints if predictions.has(\"pred_keypoints\") else None\n",
    "\n",
    "    if predictions.has(\"pred_masks\"):\n",
    "        masks = np.asarray(predictions.pred_masks)\n",
    "        masks = [GenericMask(x, self.output.height, self.output.width) for x in masks]\n",
    "    else:\n",
    "        masks = None\n",
    "\n",
    "\n",
    "    if self._instance_mode == ColorMode.SEGMENTATION and self.metadata.get(\"thing_colors\"):\n",
    "        # colors = [\n",
    "        #     self._jitter([x / 255 for x in self.metadata.thing_colors[c]]) for c in classes\n",
    "        # ]\n",
    "        colors = [\n",
    "            [x / 255 for x in self.metadata.thing_colors[c]] for c in classes\n",
    "        ]\n",
    "        alpha = 0.7\n",
    "    else:\n",
    "        colors = None\n",
    "        alpha = 0.5\n",
    "\n",
    "    if self._instance_mode == ColorMode.IMAGE_BW:\n",
    "        self.output.reset_image(\n",
    "            self._create_grayscale_image(\n",
    "                (predictions.pred_masks.any(dim=0) > 0).numpy()\n",
    "                if predictions.has(\"pred_masks\")\n",
    "                else None\n",
    "            )\n",
    "        )\n",
    "        alpha = 0.3\n",
    "\n",
    "    self.overlay_instances(\n",
    "        masks=masks,\n",
    "        boxes=boxes,\n",
    "        labels=labels,\n",
    "        keypoints=keypoints,\n",
    "        assigned_colors=colors,\n",
    "        alpha=alpha,\n",
    "    )\n",
    "    return self.output\n",
    "\n",
    "idx = -1\n",
    "# plt.imshow(\n",
    "#     show_instance_image(\n",
    "#         scene_obs['images'][idx],\n",
    "#         scene_obs['instance_2ds'][idx],\n",
    "#         scene_obs['instance_classes'][idx],\n",
    "#         segmenter.seg_id_to_name\n",
    "#     )\n",
    "# )\n",
    "\n",
    "plt.imshow(\n",
    "    show_instance_image(\n",
    "        scene_obs['images'][idx],\n",
    "        scene_obs['instance_2ds'][idx],\n",
    "        None,\n",
    "        None,\n",
    "        scale=0.7,\n",
    "        custom_class_colors = colors[:256],\n",
    "        custom_draw_fn = draw_instance_predictions,\n",
    "    )\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Voxel + Instances map\n",
    "Not working? Try building up to the full SparseVoxelMap by running the \"Debugging\" subsections below this set of cells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from home_robot.mapping.voxel import SparseVoxelMap\n",
    "from home_robot.perception.detection.detic.detic_perception import DeticPerception\n",
    "from home_robot.perception.constants import RearrangeDETICCategories\n",
    "from home_robot.utils.config import get_config\n",
    "\n",
    "from home_robot.agent.ovmm_agent.ovmm_perception import (\n",
    "    OvmmPerception,\n",
    "    build_vocab_from_category_map,\n",
    "    read_category_map_file,\n",
    ")\n",
    "\n",
    "from enum import IntEnum, auto\n",
    "\n",
    "class SemanticVocab(IntEnum):\n",
    "    FULL = auto()\n",
    "    SIMPLE = auto()\n",
    "    ALL = auto()\n",
    "\n",
    "config = get_config('configs/model/semantic_sensor/grounded_sam_perception_home_robot.yaml')[0]\n",
    "\n",
    "segmenter = OvmmPerception(\n",
    "    config=config,\n",
    "    gpu_device_id=0,\n",
    ")\n",
    "simple_vocab = RearrangeDETICCategories(dict(zip(\n",
    "    data.METAINFO['CLASS_IDS'], # IDs [1, 3, 4, 5, ..., 65]\n",
    "    data.METAINFO['CLASS_NAMES'] # [wall, floor, cabinet, ...]\n",
    ")))\n",
    "\n",
    "segmenter.update_vocabulary_list(simple_vocab, SemanticVocab.SIMPLE)\n",
    "segmenter.set_vocabulary(SemanticVocab.SIMPLE)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "p = Path('/private/home/ssax/home-robot/src/home_robot/home_robot/datasets/scannet/data/scannet_2d_instance_data/scene0000_00/instance-filt/0.png')\n",
    "pil_image = Image.open(p)\n",
    "image = np.array(pil_image, dtype=\"uint8\")  # shape is (h, w) or (h, w, 3 or 4)\n",
    "image = torch.from_numpy(image)\n",
    "image.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get detections and move to CUDA\n",
    "segmenter._segmentation.box_threshold = 0.25\n",
    "segmenter._segmentation.text_threshold = 0.25\n",
    "\n",
    "from home_robot.core.interfaces import Observations\n",
    "import dataclasses\n",
    "scene_obs['instance_map'] = []\n",
    "scene_obs['instance_classes'] = []\n",
    "scene_obs['instance_scores'] = []\n",
    "device = 'cuda'\n",
    "with torch.no_grad():\n",
    "    # Ignore_classes\n",
    "    instance_map, instance_classes, instance_scores = [], [], []\n",
    "    semantic_frames = [\n",
    "        \n",
    "    ]\n",
    "    semantic_map = []\n",
    "    \n",
    "    for im in scene_obs['images'].cpu().numpy():\n",
    "        obs = Observations(rgb = (im * 255).astype(np.uint8), gps=None, compass=None, depth=None)\n",
    "        res = segmenter.predict(obs)\n",
    "        semantic_map.append(torch.from_numpy(res.semantic).int().to(device))\n",
    "        res = res.task_observations\n",
    "        instance_map.append(torch.from_numpy(res['instance_map']).int().to(device))\n",
    "        instance_classes.append(torch.from_numpy(res['instance_classes']).int().to(device))\n",
    "        instance_scores.append(torch.from_numpy(res['instance_scores']).float().to(device))\n",
    "        semantic_frame = res['semantic_frame'] if res['semantic_frame'] is not None else res['instance_map']\n",
    "        semantic_frames.append(torch.from_numpy(semantic_frame))\n",
    "    scene_obs['instance_map'] = torch.stack(instance_map, dim=0)\n",
    "    scene_obs['instance_classes'] = instance_classes\n",
    "    scene_obs['instance_scores'] = instance_scores\n",
    "    scene_obs['images'] = scene_obs['images'].to(device)\n",
    "    scene_obs['depths'] = scene_obs['depths'].to(device)\n",
    "    scene_obs['intrinsics'] = scene_obs['intrinsics'].to(device)\n",
    "    scene_obs['poses'] = scene_obs['poses'].to(device)\n",
    "    scene_obs['semantic_frame'] = semantic_frames\n",
    "\n",
    "# Some inforamtion about the last image\n",
    "print(\"Information about the results:\")\n",
    "for k, v in res.items():\n",
    "    if v is not None:\n",
    "        print(k, v.shape, v.min(), v.max(), len(np.unique(v)))\n",
    "\n",
    "# keep_range = [-2, len(scene_obs['images'])]\n",
    "# for key in ['instance_map', 'instance_classes', 'instance_scores', 'images', 'depths', 'intrinsics', 'poses', 'semantic_frame']:\n",
    "#     scene_obs[key] = scene_obs[key][keep_range[0]:keep_range[1]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(scene_obs['images'][idx].cpu())\n",
    "plt.show()\n",
    "plt.imshow(scene_obs['depths'][idx].cpu())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.imshow((scene_obs['images'].cpu().numpy()[-1] * 255).astype(np.uint8))\n",
    "plt.imshow((scene_obs['instance_map'].cpu().numpy()[-1]).astype(np.uint8))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from home_robot.perception.encoders import ClipEncoder\n",
    "encoder = ClipEncoder(version=\"ViT-B/32\", device = 'cuda:1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# encoded_classes = [encoder.encode_text]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from home_robot.mapping.instance.instance_map import ViewMatchingConfig\n",
    "from home_robot.mapping.instance.matching import Bbox3dOverlapMethodEnum\n",
    "# Create SparseVoxelMap from observations\n",
    "svm = SparseVoxelMap(\n",
    "    background_instance_label=-1,\n",
    "    instance_memory_kwargs = dict(\n",
    "        erode_mask_num_pix=5,\n",
    "        erode_mask_num_iter=5,\n",
    "        # iou_threshold=0.3,\n",
    "        # instance_association='bbox_one_sided_iou'\n",
    "        view_matching_config = ViewMatchingConfig(\n",
    "            within_class = False,\n",
    "            box_match_mode = Bbox3dOverlapMethodEnum.ONE_SIDED_IOU,\n",
    "            box_min_iou_thresh = 0.6,\n",
    "            box_overlap_weight = 0.3,\n",
    "            visual_similarity_weight = 1.0\n",
    "        ) \n",
    "    ),\n",
    "    encoder = encoder\n",
    "    # min_iou=0.25\n",
    ")\n",
    "\n",
    "n_frames = len(scene_obs['images'])\n",
    "with torch.no_grad():\n",
    "    for i in tqdm(range(n_frames)):\n",
    "        # if i == 0: continue\n",
    "        svm.add(\n",
    "            rgb = scene_obs['images'][i],\n",
    "            depth = scene_obs['depths'][i].squeeze(-1),\n",
    "            feats = scene_obs['images'][i],\n",
    "            camera_K = scene_obs['intrinsics'][i][:3,:3],\n",
    "            camera_pose = scene_obs['poses'][i], #scene_obs['axis_align_mats'][i] @ scene_obs['poses'][i],\n",
    "            instance_image = scene_obs['instance_map'][i] ,\n",
    "            instance_classes = scene_obs['instance_classes'][i],\n",
    "            instance_scores = scene_obs['instance_scores'][i],\n",
    "            # obs_info = extra information about this timestep\n",
    "            # instance_info = list of length num_instances\n",
    "        )\n",
    "print(f\"Pre-NMS found {len(svm.get_instances())} global instances\")\n",
    "\n",
    "\n",
    "# svm.instances.global_instance_nms(0, within_category=False, nms_iou_thresh=0.3)\n",
    "print(f\"Found {len(svm.get_instances())} global instances\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "svm.show(backend='pytorch3d',\n",
    "    height=1000,\n",
    "    boxes_plot_together=False,\n",
    "    boxes_name_int_to_display_name_dict=segmenter.seg_id_to_name\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "svm.get_instances()[2]._show_point_cloud_pytorch3d()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "svm.get_instances()[6]._show_instance_view_frames(mask_out_opacity=0.9, imsize=4, n_col=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for bounds in zip(gt_bounds):\n",
    "    torch.set_printoptions(sci_mode=False)\n",
    "    print(bounds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from evaluation.obj_det import eval_bboxes_and_print\n",
    "\n",
    "# Predicted\n",
    "pred_bounds = torch.stack([inst.bounds.cpu() for inst in svm.get_instances()])\n",
    "pred_class = torch.stack([inst.category_id.cpu() for inst in svm.get_instances()])\n",
    "pred_scores = torch.stack([torch.max(torch.stack([v.score for v in ins.instance_views])).cpu() for ins in svm.get_instances()])\n",
    "\n",
    "# GT\n",
    "gt_bounds = scene_obs['boxes_aligned']\n",
    "# Map GT class labels to match the segmenter since the segmenter might have extra classes (i.e. 'background, other')\n",
    "index = torch.tensor(list(segmenter.vocabulary_id_to_name.keys()), dtype=torch.int64)\n",
    "src = torch.tensor([segmenter.name_to_seg_id[name] for name in segmenter.vocabulary_id_to_name.values()], dtype=torch.int64)\n",
    "data_idx_to_seg_idx = torch.zeros(max(segmenter.vocabulary_id_to_name.keys()) + 1, dtype=torch.int64)\n",
    "data_idx_to_seg_idx.scatter_(dim=0, index=index, src=src).int()\n",
    "gt_class = data_idx_to_seg_idx[scene_obs['box_classes']]\n",
    "\n",
    "\n",
    "_ = eval_bboxes_and_print(\n",
    "    box_gt_bounds = [gt_bounds],\n",
    "    box_gt_class = [gt_class],\n",
    "    box_pred_bounds = [pred_bounds],\n",
    "    box_pred_class = [pred_class],\n",
    "    box_pred_scores = [pred_scores],\n",
    "    match_within_class = True,\n",
    "    iou_thr = (0.25, 0.5, 0.75),\n",
    "    label_to_cat = segmenter.seg_id_to_name \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize an instance\n",
    "idx = [int(c) for c in classes].index(0)\n",
    "\n",
    "instance = svm.get_instances()[idx]\n",
    "fig = instance._show_point_cloud_pytorch3d()\n",
    "fig.show()\n",
    "_ = instance._show_instance_view_frames()\n",
    "# plt.imshow(scene_obs['semantic_frame'][n_frames - 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from eval.utils import eval_instance_bbox_detection\n",
    "results = eval_instance_bbox_detection(\n",
    "    \n",
    ")\n",
    "def eval_bboxes_and_print(\n",
    "    box_gt_bounds: Tensor,\n",
    "    box_gt_class: Tensor,\n",
    "    box_pred_bounds: Tensor,\n",
    "    box_pred_class: Tensor,\n",
    "    box_pred_scores: Tensor,\n",
    "    match_within_class: bool = True,\n",
    "    iou_thr: Sequence[float] = (0.25, 0.5, 0.75)\n",
    "):"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Debugging: Build VoxelizedPointcloud"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build voxelized pointcloud from GT points\n",
    "\n",
    "from torch_geometric.nn.pool.voxel_grid import voxel_grid\n",
    "from torch_geometric.nn.pool.consecutive import consecutive_cluster\n",
    "from torch_geometric.utils import add_self_loops, scatter\n",
    "import time\n",
    "\n",
    "from home_robot.utils.voxel import (\n",
    "    VoxelizedPointcloud,\n",
    "    voxelize, \n",
    "    reduce_pointcloud\n",
    ")\n",
    "\n",
    "# Move to CUDA \n",
    "points = pointcloud_aligned.points_packed().cuda()\n",
    "features = pointcloud_aligned.features_packed().cuda()\n",
    "device = 'cuda:0'\n",
    "for key in ['images', 'depths', 'intrinsics', 'poses']:\n",
    "    scene_obs[key] = scene_obs[key].to(device)\n",
    "\n",
    "start = time.time()\n",
    "with torch.no_grad():\n",
    "    cluster_idx, cluster_consecutive_idx, _ = voxelize(points, 0.05)\n",
    "    points_reduced, feats_reduced, weights_reduced, rgb_cluster = reduce_pointcloud(cluster_consecutive_idx, points, features, rgbs=features)\n",
    " \n",
    "end = time.time()\n",
    "print(end - start)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build voxelized pointcloud from observations\n",
    "\n",
    "from pytorch3d.structures import Pointclouds as p3dPointclouds\n",
    "from pytorch3d.vis.plotly_vis import AxisArgs\n",
    "\n",
    "from home_robot.utils.bboxes_3d_plotly import plot_scene_with_bboxes\n",
    "from home_robot.utils.voxel import VoxelizedPointcloud\n",
    "\n",
    "# Build voxel map\n",
    "min_depth = 0.1\n",
    "max_depth = 4.0\n",
    "with torch.no_grad():\n",
    "    vpc = VoxelizedPointcloud()\n",
    "\n",
    "    for i in tqdm(range(len(scene_obs['images']))):\n",
    "        rgb = scene_obs['images'][i]\n",
    "        depth = scene_obs['depths'][i].squeeze(-1)\n",
    "        camera_pose = scene_obs['poses'][i]\n",
    "        camera_K = scene_obs['intrinsics'][i]\n",
    "        \n",
    "        full_world_xyz = unproject_masked_depth_to_xyz_coordinates( # Batchable!\n",
    "            depth=depth.unsqueeze(0).unsqueeze(1),\n",
    "            pose=camera_pose.unsqueeze(0),  \n",
    "            inv_intrinsics=torch.linalg.inv(camera_K[:3, :3]).unsqueeze(0),\n",
    "        )\n",
    "\n",
    "        valid_depth = torch.full_like(rgb[:,0], fill_value=True, dtype=torch.bool)\n",
    "        if depth is not None:\n",
    "            valid_depth = (depth > min_depth) & (depth < max_depth)\n",
    "        \n",
    "        rgb = rgb[valid_depth].reshape(-1, 3)\n",
    "        world_xyz = full_world_xyz[valid_depth.flatten()]\n",
    "        \n",
    "        vpc.add(world_xyz, features=rgb, rgb=rgb)\n",
    "\n",
    "points_reduced, features_reduced, weights_reduced, rgb_reduced = vpc.get_pointcloud()\n",
    "\n",
    "# Visualize\n",
    "ptc = p3dPointclouds(\n",
    "    points = [points_reduced.detach().cpu()],\n",
    "    features = [rgb_reduced.detach().cpu()],\n",
    ")\n",
    "fig = plot_scene_with_bboxes(\n",
    "    plots = { f\"{scene_id}\": { \n",
    "                                \"Points\": ptc,\n",
    "                                \"GT points\": pointcloud_aligned,\n",
    "\n",
    "                            }\n",
    "    },\n",
    "    xaxis={\"backgroundcolor\":\"rgb(200, 200, 230)\"},\n",
    "    yaxis={\"backgroundcolor\":\"rgb(230, 200, 200)\"},\n",
    "    zaxis={\"backgroundcolor\":\"rgb(200, 230, 200)\"}, \n",
    "    axis_args=AxisArgs(showgrid=True),\n",
    "    pointcloud_marker_size=3,\n",
    "    pointcloud_max_points=200_000,\n",
    "    boxes_wireframe_width=3,\n",
    "    boxes_add_cross_face_bars=False,\n",
    "    boxes_name_int_to_display_name_dict = dict(zip([int(i) for i in data.METAINFO['seg_valid_class_ids']], data.METAINFO['classes'])),\n",
    "    boxes_plot_together=True,\n",
    "    height=1000,\n",
    "    # width=1000,\n",
    ")\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experimental: Conceptfusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gradslam.slam.pointfusion import PointFusion\n",
    "from gradslam.structures.pointclouds import Pointclouds as gsPointclouds\n",
    "from gradslam.structures.rgbdimages import RGBDImages\n",
    "\n",
    "device = 'cuda:0'\n",
    "slam = PointFusion(odom=\"icp\", dsratio=2, dist_thresh=0.01, device=device, use_embeddings=True)\n",
    "for key in ['images', 'depths', 'intrinsics', 'poses']:\n",
    "    scene_obs[key] = scene_obs[key].to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame_cur, frame_prev = None, None\n",
    "pointclouds = gsPointclouds(\n",
    "    device=device,\n",
    ")\n",
    "with torch.no_grad():\n",
    "    for i in tqdm(range(len(scene_obs['images']))):\n",
    "        frame_cur = RGBDImages(\n",
    "                scene_obs['images'][i].unsqueeze(0).unsqueeze(0),\n",
    "                scene_obs['depths'][i].unsqueeze(0).unsqueeze(0),\n",
    "                scene_obs['intrinsics'][i].unsqueeze(0).unsqueeze(0),\n",
    "                scene_obs['poses'][i].unsqueeze(0).unsqueeze(0),\n",
    "                embeddings=scene_obs['images'][i].unsqueeze(0).unsqueeze(0).half(),\n",
    "            )\n",
    "        pointclouds, _ = slam.step(pointclouds, frame_cur, frame_prev) # last arg is frame_prev but not used with odom=gt\n",
    "        frame_prev = frame_cur\n",
    "\n",
    "        # # Ignore_classes\n",
    "        # instance_map, instance_classes, instance_scores = [], [], []\n",
    "        # for im in scene_obs['images'].numpy():\n",
    "        #     obs = Observations(rgb = im * 255, gps=None, compass=None, depth=None)\n",
    "        #     res = segmenter.predict(obs).task_observations\n",
    "        #     instance_map.append(torch.from_numpy(res['instance_map']).int())\n",
    "        #     instance_classes.append(torch.from_numpy(res['instance_classes']).int())\n",
    "        #     instance_scores.append(torch.from_numpy(res['instance_scores']).float())\n",
    "        # scene_obs['instance_map'] = instance_map\n",
    "        # scene_obs['instance_classes'] = instance_classes\n",
    "        # scene_obs['instance_scores'] = instance_scores    \n",
    "        #     frame_cur = RGBDImages(\n",
    "        #         _color.unsqueeze(0).unsqueeze(0),\n",
    "        #         _depth.unsqueeze(0).unsqueeze(0),\n",
    "        #         intrinsics.unsqueeze(0).unsqueeze(0),\n",
    "        #         _pose.unsqueeze(0).unsqueeze(0),\n",
    "        #         embeddings=_embedding.unsqueeze(0).unsqueeze(0),\n",
    "        #     )\n",
    "        #     pointclouds, _ = slam.step(pointclouds, frame_cur, frame_prev)\n",
    "        #     # frame_prev = frame_cur\n",
    "        #     torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from home_robot.utils.bboxes_3d_plotly import plot_scene_with_bboxes\n",
    "from pytorch3d.structures import Pointclouds as p3dPointclouds\n",
    "from pytorch3d.vis.plotly_vis import AxisArgs\n",
    "\n",
    "ptc = p3dPointclouds(\n",
    "    points = pointclouds.points_list,\n",
    "    features = pointclouds.embeddings_padded,\n",
    ")\n",
    "fig = plot_scene_with_bboxes(\n",
    "    plots = { f\"{scene_id}\": { \n",
    "                                \"Points\": ptc,\n",
    "                                # # \"Boxes\": join_boxes_as_scene(svm.instance_bboxes3d),\n",
    "                                # \"All boxes\": global_boxes,\n",
    "                                # \"Global boxes\": global_boxes,\n",
    "                                # \"GT boxes\": gt_boxes,\n",
    "                                \"GT points\": p3dPointclouds(points=[aligned_verts[:, :3]]),\n",
    "                                # \"cameras\": cameras,\n",
    "                            }\n",
    "    },\n",
    "    xaxis={\"backgroundcolor\":\"rgb(200, 200, 230)\"},\n",
    "    yaxis={\"backgroundcolor\":\"rgb(230, 200, 200)\"},\n",
    "    zaxis={\"backgroundcolor\":\"rgb(200, 230, 200)\"}, \n",
    "    axis_args=AxisArgs(showgrid=True),\n",
    "    pointcloud_marker_size=3,\n",
    "    pointcloud_max_points=30_000,\n",
    "    boxes_wireframe_width=3,\n",
    "    boxes_add_cross_face_bars=False,\n",
    "    boxes_name_int_to_display_name_dict = dict(zip([int(i) for i in data.METAINFO['seg_valid_class_ids']], data.METAINFO['classes'])),\n",
    "    boxes_plot_together=True,\n",
    "    height=1000,\n",
    "    # width=1000,\n",
    ")\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# K = scene_obs['intrinsics'][0][:3,:3]\n",
    "# depth = scene_obs['depths'][0].squeeze().unsqueeze(0).unsqueeze(1)\n",
    "# valid_depth  = (0.1 < depth) & (depth < 4.0)\n",
    "\n",
    "# xyz = unproject_masked_depth_to_xyz_coordinates(\n",
    "#     depth = depth,\n",
    "#     mask  = ~valid_depth,\n",
    "#     pose  = torch.eye(4).unsqueeze(0),\n",
    "#     inv_intrinsics = torch.linalg.inv(K).unsqueeze(0),\n",
    "# )\n",
    "# rgb = scene_obs['images'][0].reshape(-1,3)[valid_depth.flatten()]\n",
    "# print(scene_obs['image_paths'][0])\n",
    "# print(f\"Proportion depth valid: {float(valid_depth.float().mean())}\")\n",
    "# print(f\"Depth min + max: {float(depth.flatten()[valid_depth.flatten()].min())}, {float(depth.flatten()[valid_depth.flatten()].max())}\")\n",
    "# print(\"These are the mins-maxes along each world axis. They should be in meters:\")\n",
    "# for i in range(3):\n",
    "#     print(f\"  {i}: ({float(xyz[:,i].min())}, {float(xyz[:,i].max())})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(res['semantic_frame'])\n",
    "plt.show()\n",
    "plt.imshow(scene_obs['instance_map'][-1] == 0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Create SparseVoxelMap from observations\n",
    "# from home_robot.mapping.voxel import SparseVoxelMap\n",
    "# svm = SparseVoxelMap()\n",
    "# with torch.no_grad():\n",
    "#     for i in tqdm(range(len(scene_obs['images']))):\n",
    "        \n",
    "#         # Unproject and add to SparseVoxelMap\n",
    "#         K = scene_obs['intrinsics'][i][:3,:3]\n",
    "#         depth = scene_obs['depths'][i].squeeze().unsqueeze(0).unsqueeze(1)\n",
    "#         valid_depth  = (0.1 < depth) & (depth < 4.0)\n",
    "\n",
    "#         xyz = unproject_masked_depth_to_xyz_coordinates(\n",
    "#             depth = depth,\n",
    "#             mask  = ~valid_depth,\n",
    "#             pose  = (scene_obs['axis_align_mats'][i] @ scene_obs['poses'][i]).unsqueeze(0), #torch.eye(4).unsqueeze(0),\n",
    "#             inv_intrinsics = torch.linalg.inv(K).unsqueeze(0),\n",
    "#         )\n",
    "#         rgb = scene_obs['images'][i].reshape(-1,3)[valid_depth.flatten()]\n",
    "#         svm.add(\n",
    "#             world_xyz = xyz,\n",
    "#             feats = rgb,\n",
    "#         )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create SparseVoxelMap from observations\n",
    "from home_robot.mapping.voxel import SparseVoxelMap, SparseVoxelMapWithInstances\n",
    "svm = SparseVoxelMapWithInstances(background_instance_id=-1, min_iou=0.25)\n",
    "with torch.no_grad():\n",
    "    for i in tqdm(range(len(scene_obs['images']))):\n",
    "        svm.add(\n",
    "            rgb = scene_obs['images'][i],\n",
    "            depth = scene_obs['depths'][i].squeeze(-1),\n",
    "            feats = scene_obs['images'][i],\n",
    "            cam_K = scene_obs['intrinsics'][i][:3,:3],\n",
    "            cam_to_world = scene_obs['axis_align_mats'][i] @ scene_obs['poses'][i],\n",
    "            instance_masks = scene_obs['instance_map'][i] ,\n",
    "            instance_classes = scene_obs['instance_classes'][i],\n",
    "            instance_scores = scene_obs['instance_scores'][i],\n",
    "            # obs_info = extra information about this timestep\n",
    "            # instance_info = list of length num_instances\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -> SparseVoxelMapWithInstanceViews.show(backend='pytorch3d')\n",
    "\n",
    "# Plot GT scene\n",
    "from home_robot.utils.bboxes_3d import BBoxes3D, join_boxes_as_scene, join_boxes_as_batch\n",
    "from home_robot.utils.bboxes_3d_plotly import plot_scene_with_bboxes\n",
    "from pytorch3d.vis.plotly_vis import AxisArgs\n",
    "from pytorch3d.structures import Pointclouds\n",
    "import seaborn as sns\n",
    "\n",
    "colors = torch.tensor(sns.color_palette(\"husl\", len(scene_obs['boxes_aligned'])))\n",
    "gt_boxes = BBoxes3D(\n",
    "    bounds = [scene_obs['boxes_aligned']],\n",
    "    # features = [colors[0].unsqueeze(0).expand(27,3)],\n",
    "    features = [colors],\n",
    "    names = [scene_obs['box_classes'].unsqueeze(-1)]\n",
    ")\n",
    "# # detected_boxes = BBoxes3D(\n",
    "# #     bounds = [torch.stack([v.bounds_3d for v in inst]) for inst in svm.instance_views],\n",
    "# #     # features = [colors],\n",
    "# #     names = [torch.stack([v.category_id for v in inst]).unsqueeze(-1) for inst in svm.instance_views]\n",
    "# # )\n",
    "detected_boxes = BBoxes3D(\n",
    "    bounds = [torch.cat([torch.stack([v.bounds_3d for v in inst]) for inst in svm.instance_views], dim=0)],\n",
    "    # features = [colors],\n",
    "    names = [torch.cat([torch.stack([v.category_id for v in inst]).unsqueeze(-1) for inst in svm.instance_views], dim=0)],\n",
    ")\n",
    "global_boxes = BBoxes3D(\n",
    "    bounds = [svm.global_bbox_bounds],\n",
    "    # features = [colors],\n",
    ")\n",
    "\n",
    "fig = plot_scene_with_bboxes(\n",
    "    plots = { f\"{scene_id}\": { \n",
    "                                \"Points\": svm.global_voxel_grid._pcl,\n",
    "                                # \"Boxes\": join_boxes_as_scene(svm.instance_bboxes3d),\n",
    "                                \"All boxes\": global_boxes,\n",
    "                                \"Global boxes\": global_boxes,\n",
    "                                \"GT boxes\": gt_boxes,\n",
    "                                \"GT points\": Pointclouds(points=[aligned_verts[:, :3]]),\n",
    "                                # \"cameras\": cameras,\n",
    "                            }\n",
    "    },\n",
    "    xaxis={\"backgroundcolor\":\"rgb(200, 200, 230)\"},\n",
    "    yaxis={\"backgroundcolor\":\"rgb(230, 200, 200)\"},\n",
    "    zaxis={\"backgroundcolor\":\"rgb(200, 230, 200)\"}, \n",
    "    axis_args=AxisArgs(showgrid=True),\n",
    "    pointcloud_marker_size=3,\n",
    "    pointcloud_max_points=30_000,\n",
    "    boxes_wireframe_width=3,\n",
    "    boxes_add_cross_face_bars=False,\n",
    "    boxes_name_int_to_display_name_dict = dict(zip([int(i) for i in data.METAINFO['seg_valid_class_ids']], data.METAINFO['classes'])),\n",
    "    boxes_plot_together=True,\n",
    "    height=1000,\n",
    "    # width=1000,\n",
    ")\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -> SparseVoxelMapWithInstanceViews.show(backend='pytorch3d')\n",
    "\n",
    "# Plot GT scene\n",
    "from home_robot.utils.bboxes_3d import BBoxes3D\n",
    "from home_robot.utils.bboxes_3d_plotly import plot_scene_with_bboxes\n",
    "from pytorch3d.vis.plotly_vis import AxisArgs\n",
    "from pytorch3d.structures import Pointclouds\n",
    "import seaborn as sns\n",
    "\n",
    "# colors = torch.tensor(sns.color_palette(\"husl\", len(scene_obs['boxes_aligned'])))\n",
    "# boxes = BBoxes3D(\n",
    "#     bounds = [scene_obs['boxes_aligned']],\n",
    "#     # features = [colors[0].unsqueeze(0).expand(27,3)],\n",
    "#     features = [colors],\n",
    "#     names = [scene_obs['box_classes'].unsqueeze(-1)]\n",
    "# )\n",
    "fig = plot_scene_with_bboxes(\n",
    "    plots = { f\"{scene_id}\": { \n",
    "                                # \"Points\": svm.global_voxel_grid._pcl, # svm._pcl,\n",
    "                                \"Boxes\": gt_boxes,\n",
    "                                \"GT points\": Pointclouds(points=[aligned_verts[:, :3]]),\n",
    "                                # \"cameras\": cameras,\n",
    "                            }\n",
    "    },\n",
    "    xaxis={\"backgroundcolor\":\"rgb(200, 200, 230)\"},\n",
    "    yaxis={\"backgroundcolor\":\"rgb(230, 200, 200)\"},\n",
    "    zaxis={\"backgroundcolor\":\"rgb(200, 230, 200)\"}, \n",
    "    axis_args=AxisArgs(showgrid=True),\n",
    "    pointcloud_marker_size=3,\n",
    "    pointcloud_max_points=30_000,\n",
    "    boxes_wireframe_width=3,\n",
    "    boxes_add_cross_face_bars=False,\n",
    "    boxes_name_int_to_display_name_dict = dict(zip([int(i) for i in data.METAINFO['seg_valid_class_ids']], data.METAINFO['classes'])),\n",
    "    height=1000,\n",
    "    # width=1000,\n",
    ")\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import open3d\n",
    "from home_robot.utils.point_cloud import numpy_to_pcd, pcd_to_numpy, show_point_cloud, get_xyz_coordinates\n",
    "from home_robot.mapping.voxel import SparseVoxelMap\n",
    "\n",
    "# class ObjectPointcloudRepresentation(object):\n",
    "#     def __init__(self, visualize_planner=False):\n",
    "#         self.started = False\n",
    "#         self.voxel_map = SparseVoxelMap(resolution=0.01)\n",
    "\n",
    "#     def step(self, rgb, depth, cam_to_world, K, visualize_map=False):\n",
    "#         \"\"\"Step the collector. Get a single observation of the world. Remove bad points, such as\n",
    "#         those from too far or too near the camera.\"\"\"\n",
    "#         # Get the camera pose and make sure this works properly\n",
    "#         camera_pose = cam_to_world\n",
    "#         # Get RGB and depth as necessary\n",
    "#         orig_rgb = rgb.clone().detach()\n",
    "#         orig_depth = depth.clone().detach()\n",
    "\n",
    "#         assert rgb.shape[:-1] == depth.squeeze(-1).shape, f\"{rgb.shape} {depth.shape}\"\n",
    "#         assert rgb.ndim == 3\n",
    "#         height, width, _ = rgb.shape\n",
    "        \n",
    "#         K = K[:3,:3]\n",
    "#         depth = depth.squeeze().unsqueeze(0).unsqueeze(1)\n",
    "#         valid_depth  = (0.1 < depth) & (depth < 4.0)\n",
    "#         xyz = get_xyz_coordinates(\n",
    "#             depth = depth,\n",
    "#             mask  = ~valid_depth,\n",
    "#             pose  = torch.eye(4).unsqueeze(0),\n",
    "#             inv_intrinsics = torch.linalg.inv(K).unsqueeze(0),\n",
    "#         ) \n",
    "        \n",
    "#         # apply depth filter\n",
    "#         depth = depth.reshape(-1)\n",
    "#         rgb = rgb.reshape(-1, 3)\n",
    "#         xyz = xyz.reshape(-1, 3)\n",
    "#         rgb = rgb[valid_depth.flatten()]\n",
    "#         # xyz = xyz[valid_depth.flatten()]\n",
    "#         # TODO: remove debug code\n",
    "#         # For now you can use this to visualize a single frame\n",
    "#         # show_point_cloud(xyz, rgb / 255, orig=np.zeros(3))\n",
    "#         self.voxel_map.add(\n",
    "#             camera_pose.numpy(),\n",
    "#             xyz.numpy(),\n",
    "#             rgb.numpy(),\n",
    "#             depth=depth.numpy(),\n",
    "#             K=K.numpy(),\n",
    "#             orig_rgb=orig_rgb.numpy(),\n",
    "#             orig_depth=orig_depth.numpy(),\n",
    "#         )\n",
    "#         if visualize_map:\n",
    "#             self.voxel_map.get_2d_map(debug=True)\n",
    "\n",
    "#     def get_2d_map(self):\n",
    "#         \"\"\"Get 2d obstacle map for low level motion planning and frontier-based exploration\"\"\"\n",
    "#         return self.voxel_map.get_2d_map()\n",
    "\n",
    "#     def show(self) -> Tuple[np.ndarray, np.ndarray]:\n",
    "#         \"\"\"Display the aggregated point cloud.\"\"\"\n",
    "\n",
    "#         # Create a combined point cloud\n",
    "#         # Do the other stuff we need\n",
    "#         pc_xyz, pc_rgb = self.voxel_map.get_data()\n",
    "#         show_point_cloud(pc_xyz, pc_rgb / 255, orig=np.zeros(3))\n",
    "#         return pc_xyz, pc_rgb\n",
    "\n",
    "import collections.abc    \n",
    "from typing import Optional, Union\n",
    "\n",
    "from home_robot.utils.point_cloud import numpy_to_pcd, pcd_to_numpy, show_point_cloud, get_xyz_coordinates\n",
    "from home_robot.mapping.voxel import SparseVoxelMap\n",
    "# Recursive d.update(u). Good for configs!\n",
    "def update(d, u):\n",
    "    # https://stackoverflow.com/questions/3232943/update-value-of-a-nested-dictionary-of-varying-depth\n",
    "    for k, v in u.items():\n",
    "        if isinstance(v, collections.abc.Mapping):\n",
    "            d[k] = update(d.get(k, {}), v)\n",
    "        else:\n",
    "            d[k] = v\n",
    "    return d\n",
    "\n",
    "class ObjectPointcloudRepresentation:\n",
    "\"\"\"Simple class to collect RGB, Depth, and Pose information for building 3d spatial-semantic\n",
    "maps for the robot. Needs to subscribe to:\n",
    "- color images\n",
    "- depth images\n",
    "- camera info\n",
    "- joint states/head camera pose\n",
    "- base pose (relative to world frame)\n",
    "\n",
    "This is an example collecting the data; not necessarily the way you should do it.\n",
    "\"\"\"\n",
    "\n",
    "# TODO(sasha): Remove dependence on robot and semantic_sensor.\n",
    "#   instead, step() is an interface with a TYPE of output.\n",
    "def __init__(self,\n",
    "        robot # obs = robot.get_observation(),\n",
    "        #  obs.rgb\n",
    "        #  obs.depth\n",
    "        #  obs.xyz (depth unprojected to camera coords) # potential convention mismatch\n",
    "        #  obs.camera_pose\n",
    "        semantic_sensor=None,\n",
    "        #  obs = semantic_sensor.predict(rgb, depth, xyz, camera_pose)\n",
    "        #  self.observations.append((camera_pose, xyz, feats, base_pose, obs, info))\n",
    "        #   world_xyz = trimesh.transform_points(xyz, camera_pose)\n",
    "        #   instances = InstanceView.create_from_observations(\n",
    "        #       world_xyz, valid_depth, obs, self._seq\n",
    "        #   )\n",
    "        visualize_planner=False\n",
    "    ):\n",
    "    self.robot = robot  # Get the connection to the ROS environment via agent\n",
    "    # Run detection here\n",
    "    self.semantic_sensor = semantic_sensor #\n",
    "    self.started = False\n",
    "    self.robot_model = HelloStretchKinematics(visualize=visualize_planner)\n",
    "    self.voxel_map = SparseVoxelMap(resolution=0.01)\n",
    "\n",
    "def step(\n",
    "        self,\n",
    "        rgb: torch.Tensor, \n",
    "        depth: torch.Tensor,\n",
    "        camera_pose: torch.Tensor, \n",
    "        base_pose_plane: Optional[torch.Tensor] ,\n",
    "        instance_bounds: Optional[torch.Tensor],\n",
    "        xyz: Optional[torch.Tensor] = None, #points in camera coordinates -- but we never use these and we filter xyz based on depth. \n",
    "        # Why not just remove the dependence, since it's can be tricky for people to create the XYZs themselves.\n",
    "        # Then we can avoid passing around the xyz everywhere, and we can backprop into depth if \n",
    "        visualize_map=False\n",
    "    ):\n",
    "    \"\"\"\n",
    "        Step the collector. Get a single observation of the world. Remove bad points, such as\n",
    "        those from too far or too near the camera.\n",
    "        \n",
    "        rgb: [h, w, 3]\n",
    "        depth: [h, w]\n",
    "        camera_pose: [4, 4] cam_to_world using openCV coordinates\n",
    "        base_pose_plane: Optional [x,y,theta] of the base position + orientation\n",
    "        instance_bounds: [K, 3, 2] of bounding boxes mins + maxes along each axis\n",
    "        \n",
    "    \"\"\"\n",
    "    # obs = self.robot.get_observation()\n",
    "\n",
    "    # rgb = obs.rgb\n",
    "    # depth = obs.depth\n",
    "    # xyz = obs.xyz\n",
    "    # camera_pose = obs.camera_pose\n",
    "    base_pose = np.array([obs.gps[0], obs.gps[1], obs.compass[0]])\n",
    "\n",
    "    # Get RGB and depth as necessary\n",
    "    orig_rgb = rgb.clone().detach()\n",
    "    orig_depth = depth.clone().detach()\n",
    "    assert rgb.shape[:-1] == depth.squeeze(-1).shape, f\"{rgb.shape} {depth.shape}\"\n",
    "    assert rgb.ndim == 3\n",
    "    height, width, _ = rgb.shape\n",
    "    \n",
    "    K = K[:3,:3]\n",
    "    depth = depth.squeeze().unsqueeze(0).unsqueeze(1)\n",
    "    valid_depth  = (0.1 < depth) & (depth < 4.0)\n",
    "    xyz = get_xyz_coordinates(\n",
    "        depth = depth,\n",
    "        mask  = ~valid_depth,\n",
    "        pose  = torch.eye(4).unsqueeze(0),\n",
    "        inv_intrinsics = torch.linalg.inv(K).unsqueeze(0),\n",
    "    ) \n",
    "    \n",
    "    # Semantic prediction\n",
    "    obs = self.semantic_sensor.predict(obs)\n",
    "\n",
    "    # TODO: remove debug code\n",
    "    # For now you can use this to visualize a single frame\n",
    "    # show_point_cloud(xyz, rgb / 255, orig=np.zeros(3))\n",
    "    self.voxel_map.add(\n",
    "        camera_pose,\n",
    "        xyz,\n",
    "        rgb,\n",
    "        depth=depth,\n",
    "        base_pose=base_pose,\n",
    "        obs=obs,\n",
    "        K=self.robot.head._ros_client.rgb_cam.K,\n",
    "    )\n",
    "    if visualize_map:\n",
    "        # Now draw 2d\n",
    "        self.voxel_map.get_2d_map(debug=True)\n",
    "\n",
    "def get_2d_map(self):\n",
    "    \"\"\"Get 2d obstacle map for low level motion planning and frontier-based exploration\"\"\"\n",
    "    return self.voxel_map.get_2d_map()\n",
    "\n",
    "def show(self, backend='open3d', **backend_kwargs) -> Tuple[np.ndarray, np.ndarray]:\n",
    "    \"\"\"Display the aggregated point cloud.\"\"\"\n",
    "    if backend == 'open3d':\n",
    "        self._show_open3d(self, **backend_kwargs)\n",
    "    elif backend == 'pytorch3d':\n",
    "        self._show_pytorch3d(self, **backend_kwargs)\n",
    "    else:\n",
    "        raise NotImplementedError(f\"Uknown backend {backend}, must be 'open3d' or 'pytorch3d\")\n",
    "\n",
    "def _show_pytorch3d(self, **plot_scene_kwargs):\n",
    "    from bboxes_3d_plotly import plot_scenes_with_bboxes\n",
    "\n",
    "    pc_xyz, pc_rgb = self.voxel_map.get_data()\n",
    "    pcl = Pointclouds(points=pc_xyz, features=pc_rgb) # if # channels is 3, then interpret as RGB with in [0-1]\n",
    "    \n",
    "    bounds_np = np.stack([instance_view.bounds for instance_view in self.voxel_map._instance_views], axis=0)\n",
    "    bounds = torch.tensor(bounds_np)\n",
    "    instance_bboxs = BBoxes3D(\n",
    "        boxes = bounds_np,\n",
    "        features = [features_cat], # if # channels is 3, then interpret as RGB with in [0-1]\n",
    "    )\n",
    "    \n",
    "    _default_plot_args = dict\n",
    "        xaxis={\"backgroundcolor\":\"rgb(200, 200, 230)\"},\n",
    "        yaxis={\"backgroundcolor\":\"rgb(230, 200, 200)\"},\n",
    "        zaxis={\"backgroundcolor\":\"rgb(200, 230, 200)\"}, \n",
    "        axis_args=AxisArgs(showgrid=True),\n",
    "        pointcloud_marker_size=3,\n",
    "        pointcloud_max_points=200_000\n",
    "    )\n",
    "    fig = plot_scene_with_bboxes(\n",
    "        plots = dict(\n",
    "            f\"Global scene\": { \n",
    "                \"Points\": pcl, \n",
    "                \"Instances\": instance_bboxs\n",
    "                # \"cameras\": cameras,\n",
    "            },\n",
    "            # Could add keyframes or instances here.\n",
    "        ),\n",
    "        **update(_default_plot_args, **plot_scene_kwargs)\n",
    "    )\n",
    "    \n",
    "def _show_open3d(self):\n",
    "    import open3d\n",
    "    # Create a combined point cloud\n",
    "    # Do the other stuff we need\n",
    "    pc_xyz, pc_rgb = self.voxel_map.get_data()\n",
    "    # TODO: easy version, just pt clouds\n",
    "    # show_point_cloud(pc_xyz, pc_rgb / 255, orig=np.zeros(3))\n",
    "\n",
    "    pcd = numpy_to_pcd(pc_xyz, pc_rgb / 255.0)\n",
    "    geoms = create_visualization_geometries(pcd=pcd, orig=np.zeros(3))\n",
    "    for instance_view in self.voxel_map._instance_views:\n",
    "        mins, maxs = instance_view.bounds\n",
    "        width, height, depth = maxs - mins\n",
    "\n",
    "        # Create a mesh to visualzie where the instances were seen\n",
    "        mesh_box = open3d.geometry.TriangleMesh.create_box(\n",
    "            width=width, height=height, depth=depth\n",
    "        )\n",
    "\n",
    "        # Get vertex array from the mesh\n",
    "        vertices = np.asarray(mesh_box.vertices)\n",
    "\n",
    "        # Translate the vertices to the desired position\n",
    "        vertices += mins\n",
    "        triangles = np.asarray(mesh_box.triangles)\n",
    "\n",
    "        # Create a wireframe mesh\n",
    "        lines = []\n",
    "        for tri in triangles:\n",
    "            lines.append([tri[0], tri[1]])\n",
    "            lines.append([tri[1], tri[2]])\n",
    "            lines.append([tri[2], tri[0]])\n",
    "\n",
    "        color = [1.0, 0.0, 0.0]  # Red color (R, G, B)\n",
    "        colors = [color for _ in range(len(lines))]\n",
    "        wireframe = open3d.geometry.LineSet(\n",
    "            points=open3d.utility.Vector3dVector(vertices),\n",
    "            lines=open3d.utility.Vector2iVector(lines),\n",
    "        )\n",
    "        # Get the colors and add to wireframe\n",
    "        wireframe.colors = open3d.utility.Vector3dVector(colors)\n",
    "        geoms.append(wireframe)\n",
    "\n",
    "    # Show the geometries of where we have explored\n",
    "    open3d.visualization.draw_geometries(geoms)\n",
    "    return pc_xyz, pc_rgb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rep = ObjectPointcloudRepresentation()\n",
    "for i in tqdm(range(len(result['images']))):\n",
    "    rep.step(\n",
    "        result['images'][i], \n",
    "        result['depths'][i], \n",
    "        result['poses'][i], \n",
    "        result['intrinsics'][i], \n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
